package com.aurora.demo;

import com.alibaba.fastjson.JSONObject;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer;
import org.apache.flink.streaming.connectors.elasticsearch7.ElasticsearchSink;
import org.apache.http.HttpHost;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Requests;
import org.elasticsearch.common.xcontent.XContentType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.UUID;

/**
 * 描述：Flink集成Elasticsearch Connector连接器快速入门运行demo
 * 实现实时数据流如何无缝地流向Elasticsearch
 *
 * @author 浅夏的猫
 * @version 1.0.0
 * @date 2024-02-13 22:25:58
 */
public class ElasticsearchSinkStreamJobQuickDemo {

    private static final Logger logger = LoggerFactory.getLogger(ElasticsearchSinkStreamJobQuickDemo.class);

    public static void main(String[] args) throws Exception {

        // 创建elasticsearch集群的httpHost连接
        HttpHost httpHost = new HttpHost("localhost", 9200, "http");

        List<HttpHost> httpHosts = new ArrayList<>();

        httpHosts.add(httpHost);

        // 创建elasticsearchSinkFunction函数对象，专门用于处理数据写入elasticsearchSink算子队列，会自动创建索引
        ElasticsearchSinkFunction<JSONObject> elasticsearchSinkFunction = new ElasticsearchSinkFunction<JSONObject>() {
            @Override
            public void process(JSONObject element, RuntimeContext runtimeContext, RequestIndexer indexer) {
                String transId = element.getString("transId");
                String tradeTime = element.getString("tradeTime");
                String index = "flink_" + tradeTime;
                logger.info("交易流水={},数据写入索引{}成功", transId, index);
                IndexRequest indexRequest = Requests.indexRequest().index(index).type("_doc").id(transId).source(element, XContentType.JSON);
                indexer.add(indexRequest);
            }
        };

        // 构建elasticsearchSink算子Builder
        ElasticsearchSink.Builder<JSONObject> esSinkBuilder = new ElasticsearchSink.Builder<>(httpHosts, elasticsearchSinkFunction);

        // 每个请求最多发送的文档数量
        esSinkBuilder.setBulkFlushMaxActions(1);

        // 每次发送请求的时间间隔
        esSinkBuilder.setBulkFlushInterval(1000);

        //构建elasticsearchSink算子
        ElasticsearchSink<JSONObject> sink = esSinkBuilder.build();

        // 自定义数据源，模拟生产环境交易接入,每秒下发一个json格式数据
        SourceFunction<JSONObject> dataSource = new SourceFunction<JSONObject>() {
            @Override
            public void run(SourceContext sourceContext) throws Exception {
                while (true) {
                    //交易流水号
                    String tradeId = UUID.randomUUID().toString();
                    //交易发生时间戳
                    long timeStamp = System.currentTimeMillis();
                    //交易发生金额
                    long tradeAmount = new Random().nextInt(1000);
                    //交易名称
                    String tradeName = "支付宝转账";

                    JSONObject dataObj = new JSONObject();
                    dataObj.put("transId", tradeId);
                    dataObj.put("timeStamp", timeStamp);
                    dataObj.put("tradeTime", dateUtil(timeStamp));
                    dataObj.put("tradeAmount", tradeAmount);
                    dataObj.put("tradeName", tradeName);

                    //模拟生产，每隔1秒生成一笔交易
                    Thread.sleep(1000);
                    logger.info("源交易流水={},原始报文={}", tradeId, dataObj.toJSONString());
                    sourceContext.collect(dataObj);
                }
            }

            @Override
            public void cancel() {

            }
        };

        // 创建运行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        // 构建数据源
        DataStreamSource<JSONObject> dataStreamSource = env.addSource(dataSource);

        // 数据源写入数据算子，进行输出到elasticsearch
        dataStreamSource.addSink(sink);

        // 执行任务
        env.execute();
    }

    /**
     * 描述：时间格式化工具类
     *
     * @param timestamp 时间戳
     * @return {@code String }
     */
    private static String dateUtil(long timestamp) {
        //时间戳加工
        timestamp = timestamp / 1000;
        // 将时间戳转换为 LocalDateTime 对象
        LocalDateTime dateTime = LocalDateTime.ofInstant(Instant.ofEpochSecond(timestamp), ZoneId.systemDefault());
        // 定义日期时间格式
        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyyMMdd");
        // 格式化日期时间对象为指定格式的字符串
        String dateTimeFormat = formatter.format(dateTime);
        return dateTimeFormat;
    }
}
